import asyncio
import traceback
import re
from datetime import datetime
from types import CodeType, FunctionType
import uuid
import json

from fastapi import APIRouter, Path, WebSocket, WebSocketDisconnect, status
from starlette.websockets import WebSocketState
from sqlalchemy import delete, select, update
from sqlalchemy.orm import selectinload

from schemas.base import R
from core.database import getDB
from core.exceptions import NormalException
from models.client import MqttClientTable
from models.script import ScriptTable, RScriptDeviceTable, RScriptMqttClientTable
from schemas.script import DeviceItem, ScriptSchema, ScriptOutSchema, ScriptTestResult, ScriptTestSchema
from core.device_monitor import DeviceModelMap
from core.client_creator import ClientMap


router = APIRouter(prefix="/script", tags=["北向接口脚本管理接口"])


@router.get(
    "/list/",
    summary="查看所有脚本的列表",
    response_model=R[list[ScriptOutSchema]]
)
def getScripts():
    with getDB() as db:
        scripts = db.scalars(select(ScriptTable).options(
            selectinload(ScriptTable.devices), selectinload(ScriptTable.mqttClients)
        )).fetchall()
        result = [ScriptOutSchema(
            **script.dict(exclude=["devices", "mqttClients", "devRelation", "mqttRelation"]),
            devices=[DeviceItem(no=dev.no, name=dev.name) for dev in script.devices],
            mqttClients=[cli.name for cli in script.mqttClients]
        ) for script in scripts]
    return R(data=result)


@router.get(
    "/",
    summary="获取脚本详细信息",
    response_model=R[ScriptSchema]
)
def getScriptDetail(name: str):
    with getDB() as db:
        script = db.scalars(select(ScriptTable).options(
            selectinload(ScriptTable.devices), selectinload(ScriptTable.mqttClients)
        ).where(ScriptTable.name == name)).one_or_none()
        if not script:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="查无该脚本，请刷新后重试"
            )
        result = ScriptSchema(
            **script.dict(exclude=["devices", "mqttClients", "devRelation", "mqttRelation"]),
            devices=[dev.no for dev in script.devices],
            mqttClients=[cli.name for cli in script.mqttClients]
        )
    return R(data=result)


@router.post(
    "/",
    summary="创建北向接口脚本",
    response_model=R
)
def createScript(data: ScriptSchema):
    if not data.name:
        raise NormalException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="脚本名称不允许为空"
        )
    if data.args == 1 and not data.devices:
        raise NormalException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="请至少绑定一台设备"
        )
    if data.mode == "interval" and not data.second and not data.minute and not data.hour:
        raise NormalException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="定时触发模式下必须设置间隔时间"
        )
    with getDB() as db:
        scriptExist = db.scalars(select(ScriptTable).where(ScriptTable.name == data.name)).one_or_none()
        if scriptExist:
            raise NormalException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="脚本已存在，请勿重复创建"
            )
        mqttClients = db.scalars(select(MqttClientTable).where(
            MqttClientTable.name.in_(data.mqttClients)
        )).fetchall()
        mqttMap = {cli.name: cli.id for cli in mqttClients}
        newScript = ScriptTable(
            **data.model_dump(),
            devRelation=[RScriptDeviceTable(deviceNo=no) for no in data.devices],
            mqttRelation=[RScriptMqttClientTable(mqttClientId=mqttMap[cli]) for cli in data.mqttClients]
        )
        db.add(newScript)
        db.commit()
    return R()


@router.put(
    "/",
    summary="修改脚本信息",
    response_model=R
)
def updateScript(name: str, data: ScriptSchema):
    print(data)
    with getDB() as db:
        scriptExist = db.scalars(select(ScriptTable).where(ScriptTable.name == name)).one_or_none()
        if not scriptExist:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="查无该脚本信息，请刷新后重试"
            )
        db.execute(update(ScriptTable).where(ScriptTable.name == name).values(
            **data.model_dump(exclude={"devices", "mqttClients", "devRelation", "mqttRelation"})
        ))
        db.execute(delete(RScriptDeviceTable).where(RScriptDeviceTable.scriptId == scriptExist.id))
        db.execute(delete(RScriptMqttClientTable).where(RScriptMqttClientTable.scriptId == scriptExist.id))
        mqttClients = db.scalars(select(MqttClientTable).where(
            MqttClientTable.name.in_(data.mqttClients)
        )).fetchall()
        mqttMap = {cli.name: cli.id for cli in mqttClients}
        db.add_all([
            RScriptDeviceTable(scriptId=scriptExist.id, deviceNo=no) for no in data.devices
        ] + [
            RScriptMqttClientTable(
                scriptId=scriptExist.id, mqttClientId=mqttMap[cli]
            ) for cli in data.mqttClients
        ])
        db.commit()
    return R()


@router.delete(
    "/",
    summary="删除脚本",
    response_model=R
)
def deleteScript(name: str):
    with getDB() as db:
        scriptExist = db.scalars(select(ScriptTable).where(ScriptTable.name == name)).one_or_none()
        if not scriptExist:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="查无该脚本信息，请刷新后重试"
            )
        db.execute(delete(RScriptDeviceTable).where(RScriptDeviceTable.scriptId == scriptExist.id))
        db.execute(delete(RScriptMqttClientTable).where(RScriptMqttClientTable.scriptId == scriptExist.id))
        db.execute(delete(ScriptTable).where(ScriptTable.id == scriptExist.id))
        db.commit()
    return R()


@router.post(
    "/test/",
    summary="测试脚本可行性",
    response_model=R
)
def testScript(data: ScriptTestSchema):
    device = DeviceModelMap.get(data.device)
    if not device:
        raise NormalException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"设备{data.device}不存在"
        )
    client = None
    if data.client:
        client = ClientMap.get(data.client)
        if not client:
            raise NormalException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"客户端{data.client}不存在"
            )
    # 替换执行函数中的change()条件，使其总为真
    _script = data.script
    reRes = re.findall(r"device\..*\.change\(\)", data.script)
    if reRes:
        _script = _script.replace(reRes[0], "True")
    # 从传入的脚本字符串中获取main函数用于执行
    script = compile(source=_script, filename="", mode="exec")  # 将字符串转换为compile对象
    _main = None
    for item in script.co_consts:    # 遍历字符串内存在的所有变量和对象
        if item and item.co_name and item.co_name == "main":
            _main = FunctionType(item, {"now": datetime.now})
    if not _main:
        raise AttributeError("脚本函数中不存在main函数作为入口")
    try:
        res = _main(device.model)
        if client:
            if res is None:
                raise ValueError("函数执行后无返回结果")
            if isinstance(res, dict) and res.get("topic") and res.get("payload"):
                client.send(**res)
            else:
                raise AttributeError("对接MQTT客户端的脚本返回值必须为字典且包含\"topic\"和\"payload\"字段")
            return R(data=ScriptTestResult())
        else:
            if res is None:
                raise ValueError("函数执行后无返回结果")
            if isinstance(res, dict):
                res = json.dumps(res)
            return R(data=ScriptTestResult(msg=str(res)))
    except Exception as e:
        # err = traceback.format_exc()
        return R(data=ScriptTestResult(code=400, msg=str(e)))


@router.websocket("/ws/")
async def wsCommunication(websocket: WebSocket):
    queue = None
    device = None
    rankTag = uuid.uuid4().hex
    await websocket.accept()
    try:
        srcData = await websocket.receive_json()
        data = ScriptTestSchema(**srcData)
        # 判断设备是否存在
        device = DeviceModelMap.get(data.device)
        if not device:
            await websocket.send_json(ScriptTestResult(
                code=404, msg=f"设备{data.device}不存在"
            ).model_dump())
            return
        # 判断客户端是否存在
        client = None
        if data.client:
            client = ClientMap.get(data.client)
            if not client:
                await websocket.send_json(ScriptTestResult(
                    code=404, msg=f"客户端{data.client}不存在"
                ).model_dump())
                return
        # 初始化脚本并加载函数
        script = compile(source=data.script, filename="", mode="exec")  # 将字符串转换为compile对象
        _main = None
        for item in script.co_consts:    # 遍历字符串内存在的所有变量和对象
            if item and isinstance(item, CodeType) and item.co_name and item.co_name == "main":
                _main = FunctionType(item, {"now": datetime.now})
        if not _main:
            await websocket.send_json(ScriptTestResult(
                code=404, msg=f"脚本函数中不存在main函数作为入口"
            ).model_dump())
            return
        # 将脚本插入设备采集的运行时中
        queue = device.insertTestScript(tag=rankTag, func=_main)["queue"]
        while True:
            data = await websocket.receive_text()
            if data == "quit" or data == "Q":
                return
            if queue and len(queue) > 0: 
                res = queue.pop(0)
            else:
                await asyncio.sleep(0)
                continue
            if client:
                if isinstance(res, Exception):
                    await websocket.send_json(ScriptTestResult(
                        code=400, msg=f"{str(type(res))[8:-2]}: {str(res)}"
                    ).model_dump())
                elif isinstance(res, dict) and res.get("topic") and res.get("payload"):
                    pRes = client.send(**res)
                    if isinstance(pRes, Exception):
                        await websocket.send_json(ScriptTestResult(
                            code=400, msg=f"{str(type(pRes))[8:-2]}: {str(pRes)}"
                        ).model_dump())
                    else:
                        await websocket.send_json(ScriptTestResult(
                            code=200, msg=f"数据已发送"
                        ).model_dump())
                    continue
                else:
                    await websocket.send_json(ScriptTestResult(
                        code=400, msg="对接MQTT客户端的脚本返回值必须为字典且包含\"topic\"和\"payload\"字段"
                    ).model_dump())
                    continue
            else:
                if res is None:
                    await websocket.send_json(ScriptTestResult(
                        code=400, msg="函数执行后无返回结果"
                    ).model_dump())
                    continue
                else:
                    if isinstance(res, Exception):
                        await websocket.send_json(ScriptTestResult(
                            code=400, msg=f"{str(type(res))[8:-2]}: {str(res)}"
                        ).model_dump())
                    else:
                        await websocket.send_json(ScriptTestResult(msg=res).model_dump())
                    continue
    except SyntaxError as e:
        await websocket.send_json(ScriptTestResult(
            code=500, msg=f"{str(type(e))[8:-2]}: {str(e)}"
        ).model_dump())
    except Exception as e:
        traceback.print_exc()
    finally:
        if device:
            device.removeTestScript(rankTag)
        return
